{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run tree.py\n",
    "%run meanvar_tree_utilities.py\n",
    "%run regression_tree.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_list = [100, 200, 400]; p = 10\n",
    "R = 0.1; alpha = 0.1; obj_coef = 0\n",
    "lb = 0; ub = 1;  sum_bound = 1; if_stoch_constr = False\n",
    "runs = 10\n",
    "\n",
    "generate_Y = generate_Y_lognormal\n",
    "cond_mean = [lambda x: np.exp(x[:, 0])/5, lambda x: x[:, 0]/5, lambda x: np.abs(x[:, 0])/5]\n",
    "cond_std = [lambda x: 1 - 0.5*((-3<=x[:, 1]) & (x[:, 1]<=-1)), lambda x: 1 - 0.5*((-1<=x[:, 1])&(x[:, 1]<=1)), lambda x: 1 - 0.5*((1<=x[:, 1])&(x[:, 1]<=3))]\n",
    "\n",
    "opt_solver = partial(solve_mean_variance, alpha = alpha, R = R, obj_coef = obj_coef, lb = lb, ub = ub, sum_bound = sum_bound, if_stoch_constr = if_stoch_constr)\n",
    "hessian_computer = partial(compute_hessian, alpha = alpha)\n",
    "active_constraint = partial(search_active_constraint,  R = R, lb = lb, ub = ub, sum_bound = sum_bound, if_stoch_constr = if_stoch_constr)\n",
    "gradient_computer = partial(compute_gradient,  alpha = alpha, R = R, obj_coef = obj_coef)\n",
    "update_step = partial(compute_update_step, R = R)\n",
    "\n",
    "time_list = {str(N): {key: np.zeros(runs) for key in [\"rf_approx_risk\", \"rf_approx_sol\",  \"rf_oracle\"]} for N in N_list}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for N in N_list:\n",
    "    print(\"N:\", N)\n",
    "    n_proposals = N; \n",
    "    mtry = p;\n",
    "    subsample_ratio = 1;\n",
    "    max_depth=100; \n",
    "    min_leaf_size=10; \n",
    "    balancedness_tol = 0.2; \n",
    "    honesty = False;\n",
    "    verbose = False; oracle = True;\n",
    "    bootstrap = True; \n",
    "    \n",
    "    X_list = [np.random.normal(size = (N, p)) for run in range(runs)]\n",
    "    Y_list = [generate_Y(X_list[run], cond_mean, cond_std) for run in range(runs)]\n",
    "\n",
    "    for run in range(runs):\n",
    "        print(\"run:\", run)\n",
    "        Y = Y_list[run]; Y_est = Y_list[run]\n",
    "        X = X_list[run]; X_est = X_list[run]; \n",
    "        \n",
    "        time1 = time.time()\n",
    "        rf_approx_risk = build_tree(Y, X, Y_est, X_est, \n",
    "                             opt_solver = opt_solver, hessian_computer = hessian_computer,\n",
    "                             gradient_computer = gradient_computer, \n",
    "                             search_active_constraint = active_constraint,\n",
    "                             compute_update_step = update_step,\n",
    "                             crit_computer = compute_crit_approx_risk, \n",
    "                             honesty = honesty, mtry = mtry,\n",
    "                             min_leaf_size = min_leaf_size, max_depth = max_depth, \n",
    "                             n_proposals = n_proposals, balancedness_tol = balancedness_tol,\n",
    "                             verbose = verbose)\n",
    "        time2 = time.time()\n",
    "        time_list[str(N)][\"rf_approx_risk\"][run] = time2 - time1\n",
    "\n",
    "        time1 = time.time()\n",
    "        rf_approx_sol = build_tree(Y, X, Y_est, X_est, \n",
    "                             opt_solver = opt_solver, hessian_computer = hessian_computer,\n",
    "                             gradient_computer = gradient_computer, \n",
    "                             search_active_constraint = active_constraint,\n",
    "                             compute_update_step = update_step,\n",
    "                             crit_computer = partial(compute_crit_approx_sol, obj_coef = obj_coef, alpha = alpha), \n",
    "                             honesty = honesty, mtry = mtry,\n",
    "                             min_leaf_size = min_leaf_size, max_depth = max_depth, \n",
    "                             n_proposals = n_proposals, balancedness_tol = balancedness_tol,\n",
    "                             verbose = verbose)\n",
    "        time2 = time.time()\n",
    "        time_list[str(N)][\"rf_approx_sol\"][run] = time2 - time1\n",
    "\n",
    "\n",
    "        time1 = time.time()\n",
    "        rf_oracle = build_tree(Y, X, Y_est, X_est, \n",
    "                             opt_solver = opt_solver, hessian_computer = hessian_computer,\n",
    "                             gradient_computer = gradient_computer, \n",
    "                             search_active_constraint = active_constraint,\n",
    "                             compute_update_step = update_step,\n",
    "                             crit_computer = partial(compute_crit_oracle, solver = opt_solver), \n",
    "                             honesty = honesty, mtry = mtry,\n",
    "                             min_leaf_size = min_leaf_size, max_depth = max_depth, \n",
    "                             n_proposals = n_proposals, balancedness_tol = balancedness_tol,\n",
    "                             verbose = verbose)\n",
    "        time2 = time.time()\n",
    "        time_list[str(N)][\"rf_oracle\"][run] = time2 - time1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle.dump(time_list, open(\"time_meanvar.pkl\", \"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
